




#include <torch/torch.h>
#include <ATen/ATen.h>
#include <torch/script.h>

#include <iostream>
#include <vector>
#include <memory>
#include <string>
#include <chrono>
#include <typeinfo>
#include <opencv2/opencv.hpp>
#include <opencv2/imgcodecs.hpp>
#include <opencv2/imgproc.hpp>
#include <opencv2/highgui.hpp>
#include <opencv2/objdetect/objdetect.hpp>

using namespace std;

// resize并保持图像比例不变
cv::Mat resize_with_ratio(cv::Mat& img)
{
    cv::Mat temImage;
    int w = img.cols;
    int h = img.rows;

    float t = 1.;
    float len = t * std::max(w, h);
    int dst_w = 224, dst_h = 224;
    cv::Mat image = cv::Mat(cv::Size(dst_w, dst_h), CV_8UC3, cv::Scalar(128,128,128));
    cv::Mat imageROI;
    if(len==w)
    {
        float ratio = (float)h/(float)w;
        cv::resize(img,temImage,cv::Size(224,224*ratio),0,0,cv::INTER_LINEAR);
        imageROI = image(cv::Rect(0, ((dst_h-224*ratio)/2), temImage.cols, temImage.rows));
        temImage.copyTo(imageROI);
    }
    else
    {
        float ratio = (float)w/(float)h;
        cv::resize(img,temImage,cv::Size(224*ratio,224),0,0,cv::INTER_LINEAR);
        imageROI = image(cv::Rect(((dst_w-224*ratio)/2), 0, temImage.cols, temImage.rows));
        temImage.copyTo(imageROI);
    }

    return image;
}

int main(int argc, const char* argv[])
{
   if (argc < 3) {
       std::cerr << "usage: torch_ext-app <path-to-exported-script-module> <path-to-image>\n";
       return -1;
   }

    torch::jit::script::Module module = torch::jit::load(argv[1]);
    //module.to(at::kCUDA);

    cv::Mat frame;
    cv::Mat image;
    cv::Mat input;


    frame = cv::imread(argv[2]);
    cv::resize(frame, image, cv::Size(224, 224));
    imshow("resized image", image);    //显示图像
    cv::cvtColor(image, input, cv::COLOR_BGR2RGB);

    // 下方的代码即将图像转化为Tensor，随后导入模型进行预测
    torch::Tensor tensor_image = torch::from_blob(input.data, {1, input.rows, input.cols, 3}, torch::kByte);
    tensor_image = tensor_image.permute({0,3,1,2});  //变更尺寸
    tensor_image = tensor_image.toType(torch::kFloat);
    tensor_image = tensor_image.div(255);
    //tensor_image = tensor_image.to(torch::kCUDA);
    // shape of tensor_image is N,C,H,W
    tensor_image[0][0].sub_(0.485).div_(0.229);  //減去均值,除以標準差
    tensor_image[0][1].sub_(0.456).div_(0.224);
    tensor_image[0][2].sub_(0.406).div_(0.225);
    torch::Tensor result = module.forward({tensor_image}).toTensor();

    // auto max_result = result.max(1, true);
    // auto max_index = std::get<1>(max_result).item<float>();
    // cout << max_index << endl;

    auto prob = result.softmax(1);
    // auto idx = prob.argmax();
    // cout << "The index is " << idx.item<float>() << endl;
    // cout << "The prob is " << prob[0][idx].item<float>() << endl;

    cout << "The top3 probs are: " << endl;
    auto top3 = prob.topk(3);
    cout << std::get<0>(top3) << endl;
    cout << std::get<1>(top3) << endl;

    cv::waitKey(0);
    return 0;
}







